package at.chrenko.tu.vs.shared;

import java.io.IOException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;

public class AESDecorator extends ChannelDecorator {

	private SecretKey secretKey;
	private Cipher crypt;
	private IvParameterSpec ivParam;

	public AESDecorator(Channel channel, SecretKey secretKey, byte[] ivParam) throws NoSuchAlgorithmException, NoSuchPaddingException {
		super(channel);
		this.secretKey = secretKey;
		this.ivParam = new IvParameterSpec(ivParam);
		
		this.crypt = Cipher.getInstance("AES/CTR/NoPadding");
	}
	
	@Override
	public void send(byte[] message) throws IOException {
		if(message == null)
			return;
		
		try {
			crypt.init(Cipher.ENCRYPT_MODE, secretKey, ivParam);
		} catch (InvalidKeyException e) {
			System.err.println(e.getMessage());
			return;
		} catch (InvalidAlgorithmParameterException e) {
			System.err.println(e.getMessage());
			return;
		}
		
		try {
			decoratedChannel.send(crypt.doFinal(message));
		} catch (IllegalBlockSizeException e) {
			System.err.println(e.getMessage());
			return;
		} catch (BadPaddingException e) {
			System.err.println(e.getMessage());
			return;
		}
	}
	
	@Override
	public byte[] receive() throws IOException {
		byte[] data = decoratedChannel.receive();
		
		if(data == null)
			return null;
		
		try {
			crypt.init(Cipher.DECRYPT_MODE, secretKey, ivParam);
		} catch (InvalidKeyException e) {
			System.err.println(e.getMessage());
			return null;
		} catch (InvalidAlgorithmParameterException e) {
			System.err.println(e.getMessage());
			return null;
		}
		
		try {
			return crypt.doFinal(data);
		} catch (IllegalBlockSizeException e) {
			System.err.println(e.getMessage());
			return null;
		} catch (BadPaddingException e) {
			System.err.println(e.getMessage());
			return null;
		}
	}
}
